import ast
import copy
from _ast import Call, Subscript, Attribute, AnnAssign, Expr, If, Break, Raise, UnaryOp, For
from ast import Assign, BinOp, Compare, Constant, FunctionDef, Load, Mod, Mult, NodeVisitor, Add, Name, Pow, \
    Return, Sub, Div, While, Dict, List

from typing import Any, Optional
import typing
from pprintast import pprintast
from llvmlite import ir
from MelodieAllegro.exteral_functions.builtin_ast_macros import builtin_macros
from MelodieAllegro.global_mgrs import id_generator
from MelodieAllegro.query.irquery import get_block_jump_targets
from MelodieAllegro.reference_count import should_increase_refcnt
from MelodieAllegro.type_scanner import is_valid_indexing, get_slice_value
from visitors import AssignmentVisitor
from MelodieAllegro.types import double, integer, PMBasicType, PMArrayType, PMStructType, PMType, string, \
    PMObjectType, PMFunctionType, PMStringType, PMHeapType
from MelodieAllegro.exteral_functions.builtin_func_types import builtins
from MelodieAllegro.macros import PMMacroGenerator
from MelodieAllegro.exteral_functions.builtin_function_interfaces import builtin_functions


class SpecialMessage():
    def __init__(self, msg: str):
        self.msg = msg


class FunctionVisitor(NodeVisitor):
    def __init__(self, module: ir.Module, types=None, arrays=None, custom_types=None, func_name_in_ir="") -> None:
        self.module = module
        self.func: Optional[ir.Function] = None
        self._builder: Optional[ir.IRBuilder] = None

        # 在IR中显示的函数名称
        assert isinstance(func_name_in_ir, str)
        self.func_name_in_ir = func_name_in_ir

        # Store local variables
        self.local_vars: typing.Dict[str, typing.Union[ir.Argument, ir.LoadInstr]] = {}

        # Store the type of variables
        self.local_var_types: typing.Dict[str, PMType] = types if types is not None else {}

        self.return_type = None
        self.return_var_type: Optional[PMType] = None

        # 已经注册的函数
        self.functions = {}

        # 用户定义的所有类型
        self.custom_types: typing.Dict[
            str, typing.Union[PMStructType, PMObjectType]] = {} if custom_types is None else custom_types

        # 当前如果是循环，在遇到break的时候，应当跳转到的基础块。
        self.after_current_loop_block: Optional[ir.Block] = None

        # 所有return都会跳转到最后才返回。
        self.ir_returns: typing.List[typing.Tuple[ir.Block, ir.LoadInstr]] = []
        self.return_block: Optional[ir.Block] = None
        self.return_var_names: typing.Set[str] = set()

        # 绑定引用计数变量的变量名。
        self.local_vars_with_ref: typing.Set[str] = set()

    def set_builder(self, builder: ir.IRBuilder):
        self._builder = builder

    @property
    def builder(self):
        assert self._builder is not None
        return self._builder

    def scan_scheme(self, node: FunctionDef):
        """
        扫描和生成函数的定义以及返回值。

        :param node:
        :return:
        """
        ir_func_args = []
        for i, arg in enumerate(node.args.args):
            arg_name = node.args.args[i].arg
            arg_type = self.get_var_type(arg_name)
            ir_func_args.append(arg_type.get_ref_type())

        fnty = ir.FunctionType(self.return_type, ir_func_args)
        self.func = ir.Function(self.module, fnty,
                                name=node.name if self.func_name_in_ir == '' else self.func_name_in_ir)

    def is_defined(self, varname):
        """

        :param varname:
        :return:
        """
        return varname in self.local_vars

    def get_var(self, varname: str) -> typing.Tuple[typing.Union[ir.LoadInstr, ir.Argument], PMType]:
        return self.local_vars[varname], self.local_var_types[varname]

    def get_var_value(self, varname: str) -> typing.Union[ir.LoadInstr, ir.Argument]:
        return self.local_vars[varname]

    def get_var_type(self, varname: str) -> Optional[PMType]:
        """
        获取变量的类型。
        :param varname:
        :return:
        """
        try:
            typ = self.local_var_types[varname]
            assert isinstance(typ, PMType), typ
            return typ
        except KeyError as e:
            raise ValueError(f"Error: Maybe variable {e} not declared!")

    def visit_UnaryOp(self, node: UnaryOp) -> Any:
        """
        Visit negative value!
        :param node:
        :return:
        """
        operand = self.visit(node.operand)
        assert isinstance(operand, (Name, Constant))
        var_val = self.eval_node_value(operand)
        if var_val.type == integer:
            neg_val = self.builder.neg(var_val)
        elif var_val.type == double:
            neg_val = self.builder.fneg(var_val)
        else:
            raise NotImplementedError(var_val.type)
        tmp_var_id = id_generator.new_id('neg-var')
        self.store_value(tmp_var_id, neg_val, PMBasicType(neg_val.type))
        return Name(id=tmp_var_id)

    def store_value(self, varname: str, value, value_type: Optional[PMType] = None):
        """
        存储变量
        :param varname:
        :param value:
        :param value_type:
        :return:
        """
        if self.local_vars.get(varname) is None and self.local_var_types.get(varname) is None:

            if value_type is None:
                assert value.type == integer or value.type == double or value.type == ir.IntType(1), value.type
                self.local_var_types[varname] = PMBasicType(value.type)
            else:
                assert isinstance(value_type, PMType), value_type
                self.local_var_types[varname] = value_type

            self.local_vars[varname] = value
        else:
            if isinstance(self.local_var_types[varname], PMType):
                t: PMType = self.local_var_types[varname]
                assert t.get_ref_type() == value.type, f"{varname} expected value: {t.get_ref_type()}, " \
                                                       f"but got {value.type}"
                self.local_vars[varname] = value
            else:
                raise TypeError(f"{self.local_var_types[varname]} must be a PMType!")

    def eval_node(self, node, int_as_float=False):
        assert isinstance(node, (Name, Constant)), node
        if isinstance(node, Name):
            return self.get_var_value(node.id), self.get_var_type(node.id)
        elif isinstance(node, Constant):
            if isinstance(node.value, int):
                if not int_as_float:
                    return integer(node.value), PMBasicType(integer)
                else:
                    return double(node.value), PMBasicType(double)
            elif isinstance(node.value, float):
                return double(node.value), PMBasicType(double)
            elif isinstance(node.value, str):
                byts = node.value.encode("utf8")
                return PMStringType().create(self.builder, byts), PMBasicType(string)
            else:
                raise NotImplementedError
        else:
            raise NotImplementedError

    def eval_node_value(self, node, int_as_float=False):
        """
        从Constant或者Name节点中获取一个值。

        :param node:
        :param int_as_float:
        :return:
        """
        assert isinstance(node, (Name, Constant)), node
        if isinstance(node, Name):
            return self.get_var_value(node.id)
        elif isinstance(node, Constant):
            if isinstance(node.value, int):
                if not int_as_float:
                    return integer(node.value)
                else:
                    return double(node.value)
            elif isinstance(node.value, float):
                return double(node.value)
            elif isinstance(node.value, str):
                byts = node.value.encode("utf8")
                return PMStringType().create(self.builder, byts)
            else:
                raise NotImplementedError
        else:
            raise NotImplementedError

    def start_visit(self, node: FunctionDef, ret_type) -> Any:
        if isinstance(ret_type, ir.Type):
            self.return_type = ret_type
        elif isinstance(ret_type, PMType):
            self.return_type = ret_type.get_ref_type()
            self.return_var_type = ret_type
        elif ret_type is None:
            self.return_type = ir.VoidType()
            self.return_var_type = PMBasicType(None)
        else:
            if issubclass(ret_type, int):
                self.return_type = integer
            elif issubclass(ret_type, float):
                self.return_type = double
            else:
                raise TypeError(ret_type)
        assert self.return_var_type is not None, ret_type
        assert isinstance(node, FunctionDef), pprintast(node)

        self.scan_scheme(node)

        return self.visit(node)

    def get_func_args(self, fcn: FunctionDef):
        assert isinstance(fcn, Call)
        arg_values = []
        arg_types = []
        for arg in fcn.args:
            tmp = self.visit(arg)
            if isinstance(tmp, Name):
                var = self.eval_node_value(tmp)
                arg_values.append(var)
                var_type = self.get_var_type(tmp.id)
                arg_types.append(var_type)
            elif isinstance(tmp, Constant):
                var = self.eval_node_value(tmp)
                arg_values.append(var)
                var_type = PMBasicType(var)
                arg_types.append(var_type)
            else:
                raise NotImplementedError

        return arg_values, arg_types

    def visit_Call(self, node: Call) -> Any:
        if isinstance(node.func, Name):
            if node.func.id == 'print':
                for arg in node.args:
                    arg_ = self.visit(arg)
                    v = self.eval_node_value(arg_)
                    if v.type == integer:
                        builtins.print_integer(self.builder, [v])
                    elif v.type == double:
                        builtins.print_double(self.builder, [v])
                    elif v.type == string:
                        builtins.print_string(self.builder, v)
                    else:
                        raise NotImplementedError
                builtins.print_lf(self.builder)
            elif builtin_macros.has_macro(node.func.id):
                macro = builtin_macros.get_macro(node.func.id)
                tmp_id = id_generator.new_id('macro-ret')
                value, typ = macro(self, self.builder, node.args, name=tmp_id)
                self.store_value(tmp_id, value, value_type=typ)
                return Name(id=tmp_id)
            elif builtin_functions.has_function(node.func.id):

                tmp_id = id_generator.new_id('builtin-func-ret-value')
                arg_values, arg_types = self.get_func_args(node)

                func = builtin_functions.get_function(node.func.id)
                ret = func(self.builder, arg_values, arg_types, name=tmp_id)
                if ret is not None:
                    value, typ = ret
                    self.store_value(tmp_id, value, value_type=typ)
                    return Name(id=tmp_id)
                else:
                    return SpecialMessage(f"builtin_functions {node.func.id} returns nothing")
            elif node.func.id == 'refcnt':
                assert len(node.args) == 1
                tmp = self.visit(node.args[0])
                assert isinstance(tmp, Name)
                arg, typ = self.get_var(tmp.id)
                assert isinstance(typ, (PMHeapType,)), typ
                tmp_id = id_generator.new_id('ref-count-value')
                self.store_value(tmp_id, typ.get_refcnt(self.builder, arg))
                return Name(id=tmp_id)
            elif node.func.id in self.custom_types:
                custom_type: PMStructType = self.custom_types[node.func.id]
                if isinstance(custom_type, PMStructType):
                    tmp_name = id_generator.new_id('new-struct')
                    stored = custom_type.create(self.builder, tmp_name)
                    self.store_value(tmp_name, stored, value_type=custom_type)
                    return Name(id=tmp_name)
                elif isinstance(custom_type, PMObjectType):
                    args = []
                    for arg in node.args:
                        args.append(self.eval_node_value(self.visit(arg)))
                    tmp_name = id_generator.new_id('new-object')
                    stored = custom_type.create(self.builder, args, tmp_name)
                    self.store_value(tmp_name, stored, value_type=custom_type)
                    return Name(id=tmp_name)
                else:
                    raise NotImplementedError(custom_type)
            elif node.func.id in self.functions:
                func_type: PMStructType = self.functions[node.func.id]
                assert isinstance(func_type, PMFunctionType)
                args, arg_types = self.get_func_args(node)
                var_name = id_generator.new_id('func-result')
                for arg, typ in zip(args, arg_types):
                    if isinstance(typ, PMHeapType):
                        typ.inc_ref(self.builder, arg)
                if node.func.id == self.func_name_in_ir:
                    stored = func_type.call_func_with_ir(self.builder, self.func, args, var_name)
                else:
                    stored = func_type.call_func(self.builder, args, var_name)
                self.store_value(var_name, stored, value_type=func_type.ret_type)
                return Name(id=var_name)
            else:
                raise NotImplementedError(node.func.id)
        elif isinstance(node.func, Attribute):
            # 直接调用类的方法，而不通过对象。
            if isinstance(node.func.value, Name) and node.func.value.id in self.custom_types:
                cls_type: PMObjectType = self.custom_types[node.func.value.id]
                assert isinstance(cls_type, PMObjectType)
                obj_name = self.visit(node.args[0])
                assert isinstance(obj_name, Name)
                obj_type: PMObjectType
                obj, obj_type = self.get_var(obj_name.id)

                tmp_var = id_generator.new_id('class_call_result', node.lineno)
                val, ret_type = cls_type.call_method_from_class(self.builder, node.func.attr, obj_type, obj,
                                                                [self.eval_node_value(self.visit(arg)) for arg in
                                                                 node.args[1:]], tmp_var)
                if val.type != ir.VoidType():
                    self.store_value(tmp_var, val, value_type=ret_type)
                    return Name(id=tmp_var)
                else:
                    return
            else:

                ret = self.visit(node.func)
                assert isinstance(node.func.value, Name), ast.dump(node.func.value)

                obj, obj_type = self.get_var(node.func.value.id)

                assert isinstance(obj_type, PMObjectType)
                func: PMFunctionType = self.eval_node_value(ret)
                owner_type: PMObjectType = func.owner
                tmp_var = id_generator.new_id('method_result', node.lineno)
                args_except_self = [self.eval_node_value(self.visit(arg)) for arg in node.args]
                args_values_and_typesexcept_self = [self.eval_node(self.visit(arg)) for arg in node.args]
                for arg_val, arg_type in args_values_and_typesexcept_self:
                    if isinstance(arg_type, PMHeapType):
                        arg_type.inc_ref(self.builder, arg_val)
                v = owner_type.call_method_by_functype(self.builder, func, obj_type, obj,
                                                       args_except_self,
                                                       name=tmp_var)
                self.store_value(tmp_var, v, value_type=func.ret_type)
                return Name(id=tmp_var)
        else:
            raise NotImplementedError

    def visit_Name(self, node: Name) -> Any:
        return node

    def visit_Constant(self, node: Constant) -> Any:
        return node

    def visit_List(self, node: List) -> Any:
        initial_length = len(node.elts)
        name = id_generator.new_id('list-value', node.lineno)
        v_types = set()
        v_type = None
        for index, element in enumerate(node.elts):
            value = self.visit(element)
            v = self.eval_node_value(value)
            v_types.add(v.type)
            if isinstance(value, Name):
                v_type = self.get_var_type(value.id)
            else:
                # TODO: get_value方法早晚要重写！
                v_type = PMBasicType(v.type)
        assert len(v_types) == 1, "The list is not homo"

        arr_type = PMArrayType(v_type)
        arr_ref = arr_type.create(self.builder, initial_length, name)

        for index, element in enumerate(node.elts):
            value = self.visit(element)
            v = self.eval_node_value(value)
            arr_type.setitem(self.builder, arr_ref, integer(index), v, bound_check=False)

        self.store_value(name, arr_ref, value_type=arr_type)
        return Name(id=name, ctx=Load())

    def visit_Attribute(self, node: Attribute) -> Any:
        if isinstance(node.value, Name):
            struct = self.get_var_type(node.value.id)
            if isinstance(struct, (PMStructType, PMObjectType)):
                obj = self.eval_node_value(node.value)
                name = id_generator.new_id('attribute-value')
                # 如果访问的是属性，或者实例
                if isinstance(struct, PMStructType) or (
                        isinstance(struct, PMObjectType) and struct.tell_attr_type(node.attr) == struct.PROP):
                    ret = struct.getattr(self.builder, obj, node.attr, name)
                    obj_attr_type = struct.get_attr_type(node.attr)
                    self.store_value(name, ret, value_type=obj_attr_type)
                    return Name(id=name)
                else:
                    struct: PMObjectType
                    try:
                        method = struct.get_method(node.attr)
                    except KeyError:
                        raise AttributeError(f"Object type {obj} has no attribute {node.attr}")
                    self.store_value(
                        new_id := id_generator.new_id('attribute-tmp', node.lineno if hasattr(node, 'lineno') else -1),
                        method,
                        value_type=method)
                    return Name(new_id)
            else:
                raise NotImplementedError
        elif isinstance(node.value, Attribute):
            # for example   `a.b.c`
            # node.value => `a.b`
            # node.attr =>  `.c`
            name_to_getattr = self.visit(node.value)
            assert isinstance(name_to_getattr, Name)
            obj_to_getattr = self.eval_node_value(name_to_getattr)
            obj_to_getattr_type: PMObjectType = self.get_var_type(name_to_getattr.id)
            tmp_name = id_generator.new_id('load-attr-chain', node.lineno)
            obj_attr_value = obj_to_getattr_type.getattr(self.builder, obj_to_getattr, node.attr, tmp_name)
            obj_attr_type: PMType = obj_to_getattr_type.get_attr_type(node.attr)
            self.store_value(tmp_name, obj_attr_value, value_type=obj_attr_type)
            return Name(id=tmp_name)
        elif isinstance(node.value, Subscript):
            subscr_value = self.visit(node.value)
            assert isinstance(subscr_value, Name)
            tmp_attr = Attribute(value=subscr_value, attr=node.attr, ctx=node.ctx)
            ret = self.visit_Attribute(tmp_attr)
            assert isinstance(ret, Name)
            attr_val, obj_to_getattr_type = self.get_var(ret.id)
            tmp_name = id_generator.new_id('load-attr-subscript-chain', node.lineno)
            self.store_value(tmp_name, attr_val, value_type=obj_to_getattr_type)
            return Name(id=tmp_name)

        else:
            raise NotImplementedError(node.value)

    def visit_Subscript(self, node: Subscript) -> Any:
        if isinstance(node.value, Name):
            assert isinstance(self.get_var_type(node.value.id),
                              (PMArrayType, PMObjectType)), f"variable '{node.value.id}' is un-index-able"
            index_target = self.local_vars[node.value.id]
            # if isinstance(node.slice, Index):
            #     if isinstance(node.slice.value, (Constant, Name)):
            if is_valid_indexing(node):
                var_name = id_generator.new_id('subscript', node.lineno)
                slice_value = get_slice_value(node)
                index_value = self.eval_node_value(slice_value)
                assert index_value.type == integer, f"Index should be an integer, but got {index_value.type}"
                arr_type: typing.Union[PMObjectType, PMArrayType] = self.get_var_type(node.value.id)
                ret, val_type = arr_type.getitem(self.builder,
                                                 index_target,
                                                 index_value,
                                                 name=var_name)
                self.store_value(var_name, ret, value_type=val_type)
                return Name(id=var_name, ctx=Load())
            else:
                raise NotImplementedError()
        elif isinstance(node.value, Attribute):
            res = self.visit_Attribute(node.value)
            assert isinstance(res, Name)
            new_subscr = copy.deepcopy(node)
            new_subscr.value = res
            return self.visit_Subscript(new_subscr)
        else:
            raise NotImplementedError()

    def bin_conversion(self, left_var, right_var):
        """
        对左右值进行类型转换。如果都是int,直接返回

        否则，全部转为float。
        :param left_var:
        :param right_var:
        :return:
        """
        if left_var.type == integer and right_var.type == integer:
            return left_var, right_var, integer
        else:
            if left_var.type == integer:
                left_var = self.builder.sitofp(left_var, double)
            if right_var.type == integer:
                right_var = self.builder.sitofp(right_var, double)
            return left_var, right_var, double

    def visit_If(self, node: If) -> Any:
        # 获取检验值
        assigned_vars = AssignmentVisitor().get_assigned_values(node)

        test_value = self.eval_node_value(self.visit(node.test))
        last_blk = self.builder.block
        # 保存进入if-else区块之前，已经定义的值
        defined_values = self.local_vars.copy()
        if len(node.orelse) != 0:
            with self.builder.if_else(test_value) as (then, or_else):

                with then as t:
                    print(t)
                    # bb_then: ir.Block = self.builder.basic_block
                    for item in node.body:
                        self.visit(item)
                    block_after_then = self.builder.block
                defined_values_in_then = self.local_vars.copy()
                self.local_vars = defined_values.copy()

                with or_else:
                    # bb_or_else: ir.Block = self.builder.basic_block
                    for item in node.orelse:
                        self.visit(item)
                    block_after_else = self.builder.block
                # bb_endif: ir.Block = self.builder.basic_block
                defined_values_in_else = self.local_vars.copy()

            # endif_block = self.builder.block
            # then_targets = get_block_jump_targets(bb_then)
            # orelse_targets = get_block_jump_targets(bb_or_else)

            then_targets = get_block_jump_targets(block_after_then)
            orelse_targets = get_block_jump_targets(block_after_else)

            then_jumped_to_endif = self.builder.block in then_targets
            orelse_jumped_to_endif = self.builder.block in orelse_targets

            self.local_vars = defined_values.copy()

            # 如果then和orelse两个基础块都没向endif跳转，就直接跳出去。
            if not (then_jumped_to_endif or orelse_jumped_to_endif):
                return
            else:
                # 获取在区块中定义的phi函数。
                block_phis: typing.Dict[str, ir.PhiInstr] = {}
                for var_name, _ in assigned_vars.items():
                    typ = self.get_var_type(var_name)
                    phi = self.builder.phi(typ.get_ref_type(), var_name)
                    self.store_value(var_name, phi)
                    block_phis[var_name] = phi

                for var_name, _ in assigned_vars.items():
                    phi = block_phis[var_name]
                    if then_jumped_to_endif:
                        phi.add_incoming(defined_values_in_then[var_name], block_after_then)
                    if orelse_jumped_to_endif:
                        phi.add_incoming(defined_values_in_else[var_name], block_after_else)
        else:
            with self.builder.if_then(test_value) as then:
                # with then:
                bb_then: ir.Block = self.builder.basic_block
                for item in node.body:
                    self.visit(item)
                bb_after_if = self.builder.block
                defined_values_in_then = self.local_vars.copy()

            then_targets = get_block_jump_targets(bb_after_if)

            then_jumped_to_endif = self.builder.block in then_targets

            self.local_vars = defined_values.copy()

            # 获取在区块中定义的phi函数。
            block_phis: typing.Dict[str, ir.PhiInstr] = {}
            for var_name, _ in assigned_vars.items():
                typ = self.get_var_type(var_name)
                phi = self.builder.phi(typ.get_ref_type(), var_name)
                self.store_value(var_name, phi)
                block_phis[var_name] = phi

            for var_name, _ in assigned_vars.items():
                phi = block_phis[var_name]
                phi.add_incoming(defined_values[var_name], last_blk)
                if then_jumped_to_endif:
                    phi.add_incoming(defined_values_in_then[var_name], bb_after_if)

    def visit_Expr(self, node: Expr) -> Any:
        return self.generic_visit(node)

    def visit_Break(self, node: Break) -> Any:
        self.builder.branch(self.after_current_loop_block)

    def visit_BinOp(self, node: BinOp) -> Any:
        pmg = PMMacroGenerator(self.builder)
        left: Name = self.visit(node.left)
        right: Name = self.visit(node.right)

        left_var = self.eval_node_value(left)
        right_var = self.eval_node_value(right)

        tmp_name = id_generator.new_id('binop-result')
        if isinstance(node.op, Pow):
            are_integers = [False, False]
            if left_var.type == integer:
                left_var = pmg.int_to_double(left_var)
                are_integers[0] = True
            if right_var.type == integer:
                right_var = pmg.int_to_double(right_var)
                are_integers[1] = True

            res = builtins.pow(self.builder, [
                left_var, right_var], name=tmp_name)
            if all(are_integers):
                res = pmg.double_to_int(res)
        else:
            funcs = {
                Add: (self.builder.fadd, self.builder.add),
                Sub: (self.builder.fsub, self.builder.sub),
                Mult: (self.builder.fmul, self.builder.mul),
                Div: (self.builder.fdiv, None),
                Mod: (self.builder.frem, None),
            }
            if left_var.type == integer and right_var.type == integer:
                print("all integer!")
                func = funcs[node.op.__class__]

                res = func[1](left_var, right_var, name=tmp_name)
            else:
                print("not all integer!", left_var.type, right_var.type)
                if left_var.type == integer:
                    left_var = self.builder.sitofp(left_var, double)
                if right_var.type == integer:
                    right_var = self.builder.sitofp(right_var, double)
                func = funcs[node.op.__class__]
                res = func[0](left_var, right_var, name=tmp_name)

        self.store_value(tmp_name, res)
        return Name(id=tmp_name, ctx=Load())

    def visit(self, node):
        """Visit a node."""
        method = 'visit_' + node.__class__.__name__
        assert hasattr(self, method), method
        visitor = getattr(self, method, self.generic_visit)
        return visitor(node)

    def generic_visit(self, node):
        """Called if no explicit visitor function exists for a node."""
        for field, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        self.visit(item)
            elif isinstance(value, ast.AST):
                self.visit(value)

    def visit_Raise(self, node: Raise) -> Any:
        assert isinstance(node.exc, Call)
        assert isinstance(node.exc.func, Name)
        name = node.exc.func.id
        args = node.exc.args
        assert len(args) == 3
        if name == 'IndexError':
            builtins.raise_index_error(self.builder,
                                       self.eval_node_value(self.visit(args[0])),
                                       self.eval_node_value(self.visit(args[1])),
                                       self.eval_node_value(self.visit(args[2])))
        else:
            raise NotImplementedError(f"Error type {name} is not supported to raise!")

    def visit_generic(self, node: typing.Union[AnnAssign], var_type, alloc):
        assert isinstance(node.target, Name)
        if isinstance(node.value, Call):
            if isinstance(node.value.func, Name) and builtin_macros.has_macro(node.value.func.id):
                macro = builtin_macros.get_macro(node.value.func.id)
                tmp_id = id_generator.new_id('macro-ret')
                value, typ = macro(self, self.builder, node.value.args, name=tmp_id)
                self.store_value(node.target.id, value, value_type=typ)
                return Name(id=node.target.id)
            else:
                args = []
                for arg in node.value.args:
                    args.append(self.eval_node_value(self.visit(arg)))
                tmp_name = id_generator.new_id('new-object')
                val = var_type.create(self.builder, args, tmp_name)
                self.builder.store(val, alloc)
                self.store_value(
                    node.target.id,
                    self.builder.load(alloc, name=node.target.id),
                    value_type=var_type
                )
                return Name(id=node.target.id)
        else:
            node_value = self.visit(node.value)
            assert isinstance(node_value, Name)
            # tmp_name = id_generator.new_id('new-object')
            val = self.get_var_value(node_value.id)
            self.builder.store(val, alloc)
            self.store_value(
                node.target.id,
                self.builder.load(alloc, name=node.target.id),
                value_type=var_type
            )
            return Name(id=node.target.id)

    def visit_AnnAssign(self, node: AnnAssign) -> Any:
        assert isinstance(node.target, Name)

        # annassign的目标不能已经定义
        assert not self.is_defined(node.target.id)
        var_type = self.local_var_types[node.target.id]
        alloc = self.builder.alloca(var_type.get_ref_type(), name=node.target.id)
        self.local_vars_with_ref.add(node.target.id)
        # 判断为泛型。如果是泛型，就要直接求值并且返回。
        if isinstance(var_type, PMObjectType) and var_type.name.find('<') != -1:
            ast_name = self.visit_generic(node, var_type, alloc)
            var_type = self.get_var_type(ast_name.id)
            value = self.get_var_value(ast_name.id)
            if isinstance(var_type, PMHeapType) and should_increase_refcnt(node):
                var_type.inc_ref(self.builder, value)
            return

            # 获取AnnAssign的右值。
        val = self.eval_node_value(self.visit(node.value))
        # 增加引用计数
        if isinstance(var_type, PMHeapType) and should_increase_refcnt(node):
            var_type.inc_ref(self.builder, val)
        self.builder.store(val, alloc)
        self.store_value(
            node.target.id, self.builder.load(alloc, name=node.target.id), value_type=var_type)

    def get_address(self, node):
        if isinstance(node, Attribute):
            if isinstance(node.value, Name):
                struct = self.get_var_type(node.value.id)
                if isinstance(struct, (PMStructType, PMObjectType)):
                    obj = self.eval_node_value(node.value)
                    name = id_generator.new_id('attribute-value')
                    # 如果访问的是属性，或者实例
                    if isinstance(struct, PMStructType) or (
                            isinstance(struct, PMObjectType) and struct.tell_attr_type(node.attr) == struct.PROP):
                        addr = struct.attr_address(self.builder, obj, node.attr, name)
                        return addr
                    else:
                        raise NotImplementedError
                else:
                    raise NotImplementedError
            elif isinstance(node.value, Subscript):
                pass
            else:
                raise NotImplementedError
        else:
            raise NotImplementedError

    def visit_Assign(self, node: Assign) -> Any:
        assert len(node.targets) == 1
        target: Name = node.targets[0]
        if isinstance(target, Name):
            assert self.is_defined(target.id), f"Variable {target.id} at line {target.lineno} is not defined!"
            var_type = self.get_var_type(target.id)
            value = self.visit(node.value)
            assert value is not None
            last_value = self.eval_node_value(target)
            last_value_type: PMHeapType = self.get_var_type(target.id)

            # 解析右值分别为Name或者Constant的情况。
            if isinstance(value, Name):
                val = self.eval_node_value(value)
                self.store_value(
                    target.id, val, value_type=var_type)
            elif isinstance(value, Constant):
                val = self.eval_node_value(value)
                # 这是因为只有数值常量的值不会提到最前面，在Assign中只有BasicType。
                bt = PMBasicType(val.type)
                self.store_value(
                    target.id, val, value_type=bt)
            else:
                raise NotImplementedError(f"{value} is not implemented")

            if isinstance(var_type, PMHeapType):
                assert isinstance(last_value_type, PMHeapType)
                if should_increase_refcnt(node):  # 由函数返回的值，无需增加引用计数。因此只有不是从函数返回的值才增加引用计数。
                    var_type.inc_ref(self.builder, val)
                last_value_type.dec_ref(self.builder, last_value, f"Variable {target.id}")

                self.local_vars_with_ref.add(target.id)

        elif isinstance(target, Subscript):
            target_value = self.visit(target.value)
            assert isinstance(target_value, Name)
            right_value = self.visit(node.value)
            assert isinstance(right_value, (Name, Constant)), right_value
            slice_value = self.visit(get_slice_value(target))
            index = self.eval_node_value(slice_value)

            var_to_set = self.eval_node_value(target_value)
            right_val = self.eval_node_value(right_value)
            typ: PMArrayType = self.get_var_type(target_value.id)
            if isinstance(typ, (PMArrayType, PMObjectType)):
                typ.setitem(self.builder, var_to_set, index, right_val)
            else:
                raise NotImplementedError
        elif isinstance(target, Attribute):
            right_value = self.visit(node.value)
            assert isinstance(right_value, (Name, Constant)), right_value
            name = self.visit(target.value)
            assert isinstance(name, Name)

            target_name = name.id
            var_type = self.get_var_type(target_name)
            assert isinstance(var_type, (PMStructType, PMObjectType))
            struct = self.eval_node_value(name)
            assert var_type.is_pointer
            right_val = self.eval_node_value(right_value)

            prop_type = var_type.get_property_type(target.attr)
            if isinstance(prop_type, PMHeapType):
                last_value = var_type.getattr(self.builder, struct, target.attr, 'get-attr-for-decref')
                var_type.setattr(self.builder, struct, target.attr, right_val)
                if should_increase_refcnt(node):
                    prop_type.inc_ref(self.builder, right_val)
                prop_type.dec_ref(self.builder, last_value, f"struct, {target.attr}")

            else:
                var_type.setattr(self.builder, struct, target.attr, right_val)

        else:
            raise NotImplementedError

    def visit_Compare(self, node: Compare) -> Name:
        left = self.visit(node.left)
        assert isinstance(left, Name)
        cmp_operators = {
            ast.Lt: '<',
            ast.Gt: '>',
            ast.LtE: "<=",
            ast.GtE: ">=",
            ast.Eq: "==",
            ast.NotEq: "!="
        }
        left_value = self.local_vars[left.id]
        for op, comp in zip(node.ops, node.comparators):
            comp: Name = self.visit(comp)
            comp_value = self.eval_node_value(comp)
            var_name = id_generator.new_id('comp_result')
            operator = cmp_operators[op.__class__]
            left_value, comp_value, comp_type = self.bin_conversion(
                left_value, comp_value)
            if comp_type == integer:
                cmp_result = self.builder.icmp_signed(
                    operator, left_value, comp_value, name=var_name)
            else:
                cmp_result = self.builder.fcmp_ordered(
                    operator, left_value, comp_value, name=var_name)
            self.store_value(var_name, cmp_result)
            return Name(var_name)

    def visit_For(self, node: For) -> Any:
        last_block = self.builder.block
        assigned_vars = AssignmentVisitor().get_assigned_values(node)
        for_entry = self.builder.append_basic_block(id_generator.new_id('for-loop-entry'))

        after_loop_block = self.func.append_basic_block(
            name=id_generator.new_id('after_for_loop'))

        # 保存外层循环跳出的基础块。当退出这个visit_For方法时，重新存储之。
        outer_scope_exit_block = self.after_current_loop_block
        # 保存退出循环时的基础块。
        # 如果遇到break时，添加一条分支指向self.after_current_loop_block.
        self.after_current_loop_block = after_loop_block

        iter_var_name: Name = self.visit(node.iter)
        iter_var = self.get_var_value(iter_var_name.id)

        self.builder.branch(for_entry)
        self.builder.position_at_end(for_entry)

        block_phis: typing.Dict[str, ir.PhiInstr] = {}

        for var_name, _ in assigned_vars.items():
            typ = self.get_var_type(var_name)
            phi = self.builder.phi(typ.get_ref_type(), var_name)
            phi.add_incoming(self.local_vars[var_name], last_block)
            self.store_value(var_name, phi)
            block_phis[var_name] = phi

        assert isinstance(node.target, Name)
        loop_var_name = node.target.id
        loop_var_type = self.get_var_type(node.target.id)

        iter_var_type: PMObjectType = self.get_var_type(iter_var_name.id)

        for_loop_entry = self.builder.block  # for循环开始的bb

        # Enter the body of loop
        # 判断是否已经迭代到stop。如果stop，就br到After-loop-block
        condition = self.builder.icmp_signed("==",
                                             iter_var_type.call_method(self.builder, '__stop__', iter_var_type,
                                                                       iter_var, [],
                                                                       id_generator.new_id('should-iter-stop')),
                                             integer(1), name=id_generator.new_id('should-not-loop-before-for'))
        with self.builder.if_then(condition):
            self.builder.branch(after_loop_block)
        # 在此进入主循环体。

        next_val = iter_var_type.call_method(self.builder, '__next__', iter_var_type, iter_var, [],
                                             id_generator.new_id('get_next_value'))
        self.store_value(loop_var_name, next_val, loop_var_type)

        for item in node.body:
            self.visit(item)

        for var_name, _ in assigned_vars.items():
            phi = block_phis[var_name]
            phi.add_incoming(self.local_vars[var_name], self.builder.basic_block)

        self.builder.branch(for_loop_entry)

        self.builder.position_at_end(after_loop_block)

        self.after_current_loop_block = outer_scope_exit_block
        for var_name, phi in block_phis.items():
            self.store_value(var_name, phi)
            block_phis[var_name] = phi

    def visit_While(self, node: While) -> Any:
        last_block = self.builder.block
        bb = self.func.append_basic_block(
            name=id_generator.new_id('while_loop'))

        block = self.func.append_basic_block(
            name=id_generator.new_id('after_while_loop'))

        # 保存外层循环跳出的基础块。当退出这个visit_While方法时，重新存储之。
        outer_scope_exit_block = self.after_current_loop_block
        # 保存退出循环时的基础块。
        # 如果遇到break时，添加一条分支指向self.after_current_loop_block.
        self.after_current_loop_block = block

        self.builder.branch(bb)
        self.builder.position_at_end(bb)

        assigned_vars = AssignmentVisitor().get_assigned_values(node)
        block_phis: typing.Dict[str, ir.PhiInstr] = {}

        for var_name, _ in assigned_vars.items():
            typ = self.get_var_type(var_name)
            phi = self.builder.phi(typ.get_ref_type(), var_name)
            phi.add_incoming(self.local_vars[var_name], last_block)
            self.store_value(var_name, phi)
            block_phis[var_name] = phi

        for item in node.body:
            self.visit(item)
        condition_var: Name = self.visit(node.test)
        for var_name, _ in assigned_vars.items():
            phi = block_phis[var_name]
            phi.add_incoming(self.local_vars[var_name], self.builder.basic_block)

        self.builder.cbranch(self.local_vars[condition_var.id], bb, block)
        self.builder.position_at_end(block)

        self.after_current_loop_block = outer_scope_exit_block

    def visit_Return(self, node: Return) -> Any:
        """
        对每一个Return语句，生成一个分支指向return基础块。

        同时记录当前的基础块以及求出返回的值，以便在return基础块中的phi函数使用。

        :param node:
        :return:
        """
        return_name: Name = self.visit(node.value)
        self.ir_returns.append((self.builder.block, self.eval_node_value(return_name)))
        if isinstance(return_name, Name):
            self.return_var_names.add(return_name.id)
        self.builder.branch(self.return_block)

    def build_decref(self):
        if len(self.return_var_names) >= 1:
            for return_var_name in self.return_var_names:
                if isinstance(self.return_var_type, PMHeapType):
                    ret_var = self.get_var_value(return_var_name)
                    self.return_var_type.inc_ref(self.builder, ret_var)
        for var_name in self.local_vars_with_ref:
            var, var_type = self.get_var(var_name)
            if isinstance(var_type, PMHeapType):
                var_type.dec_ref(self.builder, var, f"Variable {var_name}")

    def visit_FunctionDef(self, node: FunctionDef) -> ir.Function:

        for i, arg in enumerate(self.func.args):
            arg_name = node.args.args[i].arg
            self.store_value(arg_name, arg)
            self.local_vars_with_ref.add(arg_name)

        block = self.func.append_basic_block(name="entry")
        self.return_block = self.func.append_basic_block(name='return')
        self.set_builder(ir.IRBuilder(block))
        for subitems in node.body:
            self.visit(subitems)

        if (not self.builder.block.is_terminated):
            self.builder.branch(self.return_block)

            # 如果返回值不为void类型，需要跳转到返回块，并且
            # 在ir_returns里面添加一个值，以便phi函数使用。
            # 1、ir_returns为空，则说明整个函数没有return ，显然结构有问题。
            # 2、如果ir_returns不为空，最后一个基础块缺乏return，说明很可能是
            if not isinstance(self.return_type, ir.VoidType):
                assert len(self.ir_returns) > 0, "No return found in a non-void function!"
                self.ir_returns.append((self.builder.block, ir.Constant(self.return_type, 0)))
        # 构建返回块
        self.builder.position_at_end(self.return_block)

        if len(self.ir_returns) >= 2:
            ret_value = self.builder.phi(self.return_type)
            for from_block, value in self.ir_returns:
                ret_value.add_incoming(value, from_block)
            self.build_decref()  # decref should be executed after phi
            self.builder.ret(ret_value)
        elif len(self.ir_returns) == 1:
            ret_value = self.ir_returns[0]
            self.build_decref()
            self.builder.ret(ret_value[1])
        else:
            self.build_decref()
            self.builder.ret_void()
        with open('temp.ll', 'w') as f:
            f.write(repr(self.module))
        return self.func
